import pandas as pd
import os
import matplotlib.pyplot as plt

PATH = r'C:\Users\Jeff Levy\Desktop\replication'

def check_folder(folder):
    """helper function for subfolder creation"""
    full_path = os.path.join(PATH, folder)
    if not os.path.exists(full_path):
        os.mkdir(full_path)
    return folder

def parse_file(fname):
    """Parse raw stata output from xtvar"""
    df = pd.read_stata(os.path.join(PATH, fname))
    cols = {'step':'step', 'D02':'IRF', 'pD02_ll':'Lower', 'pD02_ul':'Upper',
            'FEVD02':'FEVD', 'pFEVD02_ll':'FEVD_Lower', 'pFEVD02_ul':'FEVD_Upper',
            'D03':'IRF_reversed', 'pD03_ll':'Lower_reversed', 'pD03_ul':'Upper_reversed'}
    df = df.rename(columns=cols)
    return df[list(cols.values())]

def irf_plots(data, subset=None):
    if subset == 'reversed':
        data = {k:d for k, d in data.items() if 'reversed' in k}
        nsubs = 3
    else:
        data = {k:d for k, d in data.items() if 'reversed' not in k and 'lag' not in k}
        nsubs = 3

    name_keys = {'hf_ln_employment':'HF log(uncertainty) -> HF log(employment)',
                 'hf_unemp_rate':'HF log(uncertainty) -> HF unemployment_rate',
                 'hp_ln_employment':'HP log(uncertainty) -> HP log(employment)',
                 'hp_unemp_rate':'HP log(uncertainty) -> HP unemployment_rate',
                 'ln_employment':'log(uncertainty) -> log(employment)',
                 'unemp_rate':'log(uncertainty) -> unemployment_rate',
                 'hp_unemp_rate_reversed':'HP log(uncertainty) -> HP unemployment_rate',
                 'unemp_rate_reversed':'log(uncertainty) -> unemployment_rate',
                 'hf_unemp_rate_reversed':'HF log(uncertainty) -> HF unemployment_rate'}

    #calculated from xtvar_ready.csv, 44 state sample (using modified xtvar.ado to expose reduced form results)
    unit_sds_u_unemp = {'log(uncertainty)':  0.45126,
                        'HP log(uncertainty)': 0.43936,
                        'HF log(uncertainty)': 0.44873}
    unit_sds_u_emp =   {'log(uncertainty)':  0.45215,
                        'HP log(uncertainty)': 0.44032,
                        'HF log(uncertainty)': 0.44962}
    unit_sds_e = {'HF log(employment)': 0.67691,
                  'HP log(employment)': 0.09833,
                  'log(employment)': 0.10108,
                  'HF unemployment_rate': 0.38243,
                  'HP unemployment_rate': 0.11305,
                  'unemployment_rate': 0.11644}

    emp = {k.replace('.dta', ''):v for k, v in data.items() if 'employment' in k}
    unemp = {k.replace('.dta', '').replace('top12_', ''):v for k, v in data.items() if 'unemp_rate' in k}

    for label, dat in [('Employment (log)', emp), ('Unemployment Rate', unemp)]:
        fig, axs = plt.subplots(nsubs, 1, figsize=(8,10), sharey=True)
        for (kind, d), ax in zip(dat.items(), axs):
            if 'reversed' in kind:
                irf, lower, upper = 'IRF_reversed', 'Lower_reversed', 'Upper_reversed'
            else:
                irf, lower, upper = 'IRF', 'Lower', 'Upper'
            x = d['step']
            ax.plot(x, d[irf], 'k-')
            ax.plot(x, d[lower], 'k--', alpha=0.5)
            ax.plot(x, d[upper], 'k--', alpha=0.5)
            ax.axhline(0, 0, 1, color='k', linestyle='-', linewidth=0.2)
            if 'Employment' in label:
                peak = d[irf].min()
                peak_pos = peak*1.14
            else:
                peak = d[irf].max()
                peak_pos = peak*1.05
            peak_date = d[d[irf] == peak].step.values[0]
            ax.annotate(str(round(peak, 3)).ljust(5, "0"), (peak_date, peak_pos), fontsize='large')
            for k in ['HF log(employment)', 'HP log(employment)', 'log(employment)', 'HF unemployment_rate', 'HP unemployment_rate', 'unemployment_rate']:
                if k in name_keys[kind]:
                    if 'Employment' in label:
                        e = 'Employment \u03C3'
                    else:
                        e = 'Unemp Rate \u03C3'
                    if subset == 'top12':
                        e_value = str(round(unit_sds_top12[k], 3)).ljust(5, "0")
                    else:
                        e_value = str(round(unit_sds_e[k], 3)).ljust(5, "0")
                    break
            for k in ['HP log(uncertainty)', 'HF log(uncertainty)', 'log(uncertainty)']:
                if k in name_keys[kind]:
                    if subset == 'top12':
                        edict = unit_sds_top12
                        unedict = unit_sds_top12
                    else:
                        edict = unit_sds_u_emp
                        unedict = unit_sds_u_unemp
                    if 'Employment' in label:
                        unit_sd_size = f'EPU-C \u03C3: {str(round(edict[k], 3)).ljust(5, "0")}'
                    else:
                        unit_sd_size = f'EPU-C \u03C3: {str(round(unedict[k], 3)).ljust(5, "0")}'
                    break
            if not subset:
                ax.text(.98, .91, unit_sd_size.ljust(3, '0'), fontsize='large', horizontalalignment='right', transform=ax.transAxes)
        axs[0].set_ylabel('Hamilton Filter', fontsize='xx-large')
        axs[1].set_ylabel('HP Filter', fontsize='xx-large')
        axs[2].set_ylabel('No Filter', fontsize='xx-large')
        for a in axs:
            a.xaxis.set_tick_params(labelsize='x-large')
            a.yaxis.set_tick_params(labelsize='x-large')

        if subset:
            out = '_'.join([label, subset]) + '.png'
        else:
            out = label + '.png'
        fig.tight_layout()
        outpath = os.path.join(PATH, check_folder('irfs'), out)
        if outpath != 'Employment (log)_reversed.png':
            fig.savefig(os.path.join(PATH, check_folder('irfs'), out))
            print('Saved:', os.path.join(PATH, 'irfs', out))
        plt.close()

def irf_plots_lagcomp(data, subset=None):
    formula = 'HF log(uncertainty) -> HF unemployment_rate'
    fig, axs = plt.subplots(3, 1, figsize=(8,10), sharey=True)
    name_keys = {'hf_unemp_rate_3lag.dta':'3 Lags', 'hf_unemp_rate.dta':'6 Lags', 'hf_unemp_rate_12lag.dta':'12 Lags'}
    for name_key, ax in zip(name_keys.keys(), axs):
        d = data[name_key]
        x = d['step']
        ax.plot(x, d['IRF'], 'k-')
        ax.plot(x, d['Lower'], 'k--', alpha=0.5)
        ax.plot(x, d['Upper'], 'k--', alpha=0.5)
        ax.axhline(0, 0, 1, color='k', linestyle='-', linewidth=0.2)
        ax.set_ylabel(name_keys[name_key], fontsize='xx-large')
        ax.xaxis.set_tick_params(labelsize='x-large')
        ax.yaxis.set_tick_params(labelsize='x-large')

        peak = d.IRF.max()
        peak_date = d[d.IRF == peak].step.values[0]
        ax.annotate(str(round(peak, 3)).ljust(5, "0"), (peak_date, peak*1.05))

    out = 'HF_lags_comp.png'
    fig.tight_layout()
    fig.savefig(os.path.join(PATH, check_folder('irfs'), out))
    print('Saved:', os.path.join(PATH, 'irfs', out))
    plt.close()

if __name__ == '__main__':
    replication_exclude = ['all_state_level_epu.dta',
                           'all_state_level_epu_wide.dta',
                           'Merged_Working_Data.dta',
                           'final_state_month_covid_npi.dta',
                           'npi_data_summer_2020_onwards.dta',
                           'npi_schools_weekly.dta',
                           'npi_school_data_monthly.dta']
    fnames = [f for f in os.listdir(PATH) if f.endswith('.dta') and f not in replication_exclude]
    data = {f:parse_file(f) for f in fnames}
    irf_plots(data)
    irf_plots(data, subset='reversed')
    irf_plots_lagcomp(data)
